import pandas as pd
import pymysql
from sqlalchemy import create_engine, text
from docx import Document
from docx.shared import Inches, Pt, RGBColor
from docx.enum.text import WD_ALIGN_PARAGRAPH
from docx.oxml.ns import qn
from docx.oxml import parse_xml
import os
from datetime import datetime
import json

# 数据库配置（请修改为你的配置）
DB_CONFIG = {
    "host": "localhost",
    "port": 3306,
    "user": "root",
    "password": "",  # 替换为你的密码
    "db": "stock_db",
    "charset": "utf8mb4"
}

# 图片类型映射 - 根据实际的数据库记录进行修正
business_keywords = {
    '对比图表': [('comparison', 'all')],
    '季K线图表': [('individual', 'quarterly')],
    '年K线图表': [('individual', 'yearly')],
    '月K线蜡烛图': [('candle', 'monthly')]
}


def create_database_connection():
    """创建数据库连接"""
    try:
        engine = create_engine(
            f"mysql+pymysql://{DB_CONFIG['user']}:{DB_CONFIG['password']}@{DB_CONFIG['host']}:{DB_CONFIG['port']}/{DB_CONFIG['db']}?charset={DB_CONFIG['charset']}"
        )
        return engine
    except Exception as e:
        print(f"数据库连接失败: {str(e)}")
        raise


def read_company_detail_data(engine, limit=None):
    """从company_detail_analysis表读取公司数据"""
    try:
        query = """
        SELECT 
            id AS '主键ID',
            stock_code AS '股票代码',
            company_name AS '上市公司名称',
            establishment_date AS '成立日期',
            manager AS '管理者',
            manager_introduction AS '管理者介绍',
            registered_address AS '注册地址',
            company_type AS '公司性质',
            primary_industry AS '一级行业',
            secondary_industry AS '二级行业',
            tertiary_industry AS '三级行业',
            business_scope AS '经营范围',
            main_business AS '主营业务',
            violation_time AS '违规时间',
            violation_type AS '违规性质',
            violation_description AS '违规描述',
            listing_date AS '上市时间',
            controlling_shareholder AS '控股股东',
            business_development_analysis AS '业务发展分析',
            financial_condition_analysis AS '财务状况分析',
            profit_model AS '盈利模式',
            moat AS '护城河',
            customer_group AS '客户群体',
            competitors AS '竞争对手',
            business_area AS '业务区域',
            subsidiaries AS '子企业',
            core_products_services AS '核心产品或服务',
            related_upstream_downstream_listed_companies AS '关联上下游上市企业',
            main_factors_affecting_operation AS '经营状况主要影响因素',
            former_names AS '企业曾用名',
            csrc_warning_letters AS '被证监会发过警告函情况',
            major_shareholder_reduction AS '大股东减持情况',
            company_goodwill AS '公司商誉',
            goodwill_scale AS '商誉规模',
            asset_provision_information AS '资产计提信息',
            position_in_subdivision_industry AS '细分行业中地位',
            debt_ratio_5years AS '近5年的负债率',
            revenue_5years AS '近5年的营业收入',
            deducted_net_profit_5years AS '近5年的扣费净利润',
            gross_profit_margin_5years AS '近5年的毛利率',
            same_industry_market_value_ranking AS '同行业市值排名公司',
            dominant_products_services_in_market AS '市场中占主导地位的产品服务',
            company_highlights AS '公司亮点',
            company_risk_points AS '公司风险点',
            themes_concepts AS '题材/概念',
            what_company_does AS '公司干什么的',
            major_customers AS '主要客户',
            stock_price_surge_periods AS '股票大涨时间段',
            performance_indicators_during_surge AS '股票大涨时间段内业绩指标详细描述',
            is_cyclical AS '是否为周期股',
            stock_price_characteristics AS '此公司的股价变动特点',
            analysis_time AS '分析时间',
            batch_id AS '批次ID',
            created_at AS '创建时间'
        FROM company_detail_analysis
        """
        if limit:
            query += f" LIMIT {limit}"

        df = pd.read_sql(query, engine)
        print(f"成功读取 {len(df)} 条公司详细数据")
        return df
    except Exception as e:
        print(f"读取数据失败: {str(e)}")
        raise


def get_stock_images_from_db(engine, stock_code):
    """从stock_chart_images表获取股票的图片路径"""
    try:
        query = text("""
        SELECT chart_type, period, image_path 
        FROM stock_chart_images 
        WHERE stock_code = :stock_code
        ORDER BY chart_type, period
        """)

        with engine.connect() as connection:
            result = connection.execute(query, {'stock_code': stock_code})
            rows = result.fetchall()

        if rows:
            df = pd.DataFrame(rows, columns=['chart_type', 'period', 'image_path'])
            return df
        else:
            return pd.DataFrame()

    except Exception as e:
        print(f"获取股票图片失败（{stock_code}）: {str(e)}")
        return pd.DataFrame()


def set_chinese_font(run, font_name='微软雅黑'):
    """设置中文字体"""
    run.font.name = font_name
    run._element.rPr.rFonts.set(qn('w:eastAsia'), font_name)


def add_company_header(doc, company_data):
    """添加公司头部信息"""
    # 主标题
    title = doc.add_heading('', level=1)
    title_run = title.add_run(f"{company_data['上市公司名称']}（{company_data['股票代码']}）分析报告")
    title_run.font.size = Pt(18)
    title_run.bold = True
    title_run.font.color.rgb = RGBColor(0, 0, 128)  # 深蓝色
    set_chinese_font(title_run)
    title.alignment = WD_ALIGN_PARAGRAPH.CENTER

    doc.add_paragraph()  # 空行

    # 基本信息表格
    table = doc.add_table(rows=6, cols=2)
    table.style = 'Light Grid Accent 1'
    table.autofit = False

    # 设置列宽
    for row in table.rows:
        row.cells[0].width = Inches(1.5)
        row.cells[1].width = Inches(4.5)

    # 设置表格基本信息
    info_data = [
        ('公司名称', company_data['上市公司名称']),
        ('股票代码', company_data['股票代码']),
        ('成立日期', format_date(company_data['成立日期'])),
        ('上市时间', format_date(company_data['上市时间'])),
        ('公司性质', format_value(company_data['公司性质'])),
        ('注册地址', format_value(company_data['注册地址']))
    ]

    for i, (label, value) in enumerate(info_data):
        row_cells = table.rows[i].cells
        row_cells[0].text = label
        row_cells[1].text = str(value)

        # 设置字体
        for cell in row_cells:
            for paragraph in cell.paragraphs:
                for run in paragraph.runs:
                    set_chinese_font(run)
                    if cell == row_cells[0]:  # 标签单元格加粗
                        run.bold = True
                    run.font.size = Pt(10.5)

    doc.add_paragraph()  # 空行


def format_value(value):
    """格式化值，处理空值和None"""
    if pd.isna(value) or value is None or str(value).strip() in ['', '暂无信息', 'None']:
        return '暂无'
    return str(value).strip()


def format_date(date_value):
    """格式化日期"""
    if pd.isna(date_value) or date_value is None or str(date_value).strip() in ['', '暂无信息', 'None']:
        return '暂无'

    try:
        if isinstance(date_value, str):
            # 尝试解析日期字符串
            if len(date_value) == 10:  # YYYY-MM-DD格式
                return date_value
            else:
                return str(date_value)
        else:
            return str(date_value)
    except:
        return str(date_value)


def add_company_overview(doc, company_data):
    """添加公司概况"""
    doc.add_heading('一、公司概况', level=1)

    # 公司基本信息
    overview_info = [
        ('一级行业', company_data['一级行业']),
        ('二级行业', company_data['二级行业']),
        ('三级行业', company_data['三级行业']),
        ('控股股东', company_data['控股股东']),
        ('企业曾用名', company_data['企业曾用名']),
        ('公司干什么的', company_data['公司干什么的'])
    ]

    for label, value in overview_info:
        formatted_value = format_value(value)
        if formatted_value != '暂无':
            p = doc.add_paragraph()
            p.add_run(f"{label}：").bold = True
            p.add_run(formatted_value)
            set_chinese_font(p.runs[0])
            set_chinese_font(p.runs[1])


def add_business_analysis(doc, company_data):
    """添加业务分析"""
    doc.add_heading('二、业务分析', level=1)

    # 经营范围
    business_scope = format_value(company_data['经营范围'])
    if business_scope != '暂无':
        doc.add_heading('## 经营范围', level=2)
        # 对经营范围进行适当分段
        scope_text = business_scope.replace('；', '；\n')
        paragraphs = scope_text.split('\n')
        for para in paragraphs:
            if para.strip():
                p = doc.add_paragraph(para.strip())
                set_chinese_font(p.runs[0])

    # 主营业务
    main_business = format_value(company_data['主营业务'])
    if main_business != '暂无':
        doc.add_heading('## 主营业务', level=2)
        p = doc.add_paragraph(main_business)
        set_chinese_font(p.runs[0])

    # 核心产品或服务
    core_products = format_value(company_data['核心产品或服务'])
    if core_products != '暂无':
        doc.add_heading('## 核心产品或服务', level=2)
        p = doc.add_paragraph(core_products)
        set_chinese_font(p.runs[0])

    # 市场中占主导地位的产品/服务
    dominant_products = format_value(company_data['市场中占主导地位的产品服务'])
    if dominant_products != '暂无':
        doc.add_heading('## 主导产品/服务', level=2)
        p = doc.add_paragraph(dominant_products)
        set_chinese_font(p.runs[0])


def add_market_analysis(doc, company_data):
    """添加市场分析"""
    doc.add_heading('三、市场分析', level=1)

    analysis_items = [
        ('## 客户群体', company_data['客户群体']),
        ('## 主要客户', company_data['主要客户']),
        ('## 业务区域', company_data['业务区域']),
        ('## 盈利模式', company_data['盈利模式']),
        ('## 护城河', company_data['护城河']),
        ('## 细分行业中地位', company_data['细分行业中地位'])
    ]

    for label, value in analysis_items:
        formatted_value = format_value(value)
        if formatted_value != '暂无':
            doc.add_heading(label, level=2)
            p = doc.add_paragraph(formatted_value)
            set_chinese_font(p.runs[0])

    doc.add_paragraph()  # 空行


def add_competitive_analysis(doc, company_data):
    """添加竞争分析"""
    doc.add_heading('四、竞争分析', level=1)

    # 竞争对手
    competitors = format_value(company_data['竞争对手'])
    if competitors != '暂无':
        doc.add_heading('## 竞争对手', level=2)
        p = doc.add_paragraph(competitors)
        set_chinese_font(p.runs[0])

    # 同行业市值排名公司
    industry_ranking = format_value(company_data['同行业市值排名公司'])
    if industry_ranking != '暂无':
        doc.add_heading('## 同行业公司对比', level=2)
        p = doc.add_paragraph(industry_ranking)
        set_chinese_font(p.runs[0])


def format_financial_data(value):
    """格式化财务数据"""
    formatted_value = format_value(value)
    if formatted_value == '暂无':
        return None

    try:
        # 尝试解析JSON格式的财务数据
        financial_data = json.loads(value)
        if isinstance(financial_data, dict):
            formatted_text = ""
            for year, metric_value in financial_data.items():
                formatted_text += f"{year}: {metric_value}\n"
            return formatted_text.strip()
        else:
            return formatted_value
    except:
        return formatted_value


def add_financial_analysis(doc, company_data):
    """添加财务分析"""
    doc.add_heading('五、财务分析', level=1)

    # 财务状况分析
    financial_analysis = format_value(company_data['财务状况分析'])
    if financial_analysis != '暂无':
        doc.add_heading('## 财务状况分析', level=2)
        p = doc.add_paragraph(financial_analysis)
        set_chinese_font(p.runs[0])

    # 财务指标数据
    financial_metrics = [
        ('## 近5年负债率', company_data['近5年的负债率']),
        ('## 近5年营业收入', company_data['近5年的营业收入']),
        ('## 近5年扣非净利润', company_data['近5年的扣费净利润']),
        ('## 近5年毛利率', company_data['近5年的毛利率'])
    ]

    for label, value in financial_metrics:
        formatted_data = format_financial_data(value)
        if formatted_data:
            doc.add_heading(label, level=2)
            p = doc.add_paragraph(formatted_data)
            set_chinese_font(p.runs[0])

    # 商誉信息
    goodwill = format_value(company_data['公司商誉'])
    goodwill_scale = format_value(company_data['商誉规模'])
    if goodwill != '暂无':
        doc.add_heading('## 公司商誉', level=2)
        goodwill_text = goodwill
        if goodwill_scale != '暂无':
            goodwill_text += f"（规模：{goodwill_scale}）"
        p = doc.add_paragraph(goodwill_text)
        set_chinese_font(p.runs[0])

    # 资产计提信息
    asset_provision = format_value(company_data['资产计提信息'])
    if asset_provision != '暂无':
        doc.add_heading('## 资产计提信息', level=2)
        p = doc.add_paragraph(asset_provision)
        set_chinese_font(p.runs[0])


def add_development_analysis(doc, company_data):
    """添加发展分析"""
    doc.add_heading('六、业务发展分析', level=1)

    # 业务发展分析
    business_development = format_value(company_data['业务发展分析'])
    if business_development != '暂无':
        p = doc.add_paragraph(business_development)
        set_chinese_font(p.runs[0])

    # 公司亮点
    company_highlights = format_value(company_data['公司亮点'])
    if company_highlights != '暂无':
        doc.add_heading('## 公司亮点', level=2)
        p = doc.add_paragraph(company_highlights)
        set_chinese_font(p.runs[0])

    # 题材/概念
    themes_concepts = format_value(company_data['题材/概念'])
    if themes_concepts != '暂无':
        doc.add_heading('## 题材概念', level=2)
        p = doc.add_paragraph(themes_concepts)
        set_chinese_font(p.runs[0])


def add_risk_analysis(doc, company_data):
    """添加风险分析"""
    doc.add_heading('七、风险分析', level=1)

    risk_items = [
        ('经营状况主要影响因素', company_data['经营状况主要影响因素']),
        ('公司风险点', company_data['公司风险点']),
        ('证监会警告函', company_data['被证监会发过警告函情况']),
        ('大股东减持', company_data['大股东减持情况']),
        ('是否为周期股', company_data['是否为周期股'])
    ]

    # 违规情况单独处理
    violation_time = format_value(company_data['违规时间'])
    violation_type = format_value(company_data['违规性质'])
    violation_desc = format_value(company_data['违规描述'])

    if violation_time != '暂无' and violation_time not in ['', 'None']:
        violation_text = f"{violation_time}"
        if violation_type != '暂无':
            violation_text += f" - {violation_type}"
        if violation_desc != '暂无':
            violation_text += f": {violation_desc}"
        risk_items.insert(2, ('违规情况', violation_text))

    has_risk_data = False
    for label, value in risk_items:
        formatted_value = format_value(value)
        if formatted_value != '暂无':
            has_risk_data = True
            p = doc.add_paragraph(style='List Bullet')
            p.add_run(f"{label}：").bold = True
            p.add_run(formatted_value)
            set_chinese_font(p.runs[0])
            set_chinese_font(p.runs[1])

    if not has_risk_data:
        p = doc.add_paragraph('暂无重大风险事件记录')
        set_chinese_font(p.runs[0])


def add_corporate_structure(doc, company_data):
    """添加公司结构"""
    doc.add_heading('八、公司结构', level=1)

    structure_items = [
        ('## 子企业', company_data['子企业']),
        ('## 关联上下游上市企业', company_data['关联上下游上市企业'])
    ]

    has_structure_data = False
    for label, value in structure_items:
        formatted_value = format_value(value)
        if formatted_value != '暂无':
            has_structure_data = True
            doc.add_heading(label, level=2)
            p = doc.add_paragraph(formatted_value)
            set_chinese_font(p.runs[0])

    if not has_structure_data:
        p = doc.add_paragraph('暂无公司结构信息')
        set_chinese_font(p.runs[0])


def add_management_info(doc, company_data):
    """添加管理层信息"""
    manager = format_value(company_data['管理者'])
    if manager != '暂无':
        doc.add_heading('九、管理层信息', level=1)

        p = doc.add_paragraph()
        p.add_run('管理者：').bold = True
        p.add_run(manager)
        set_chinese_font(p.runs[0])
        set_chinese_font(p.runs[1])

        manager_intro = format_value(company_data['管理者介绍'])
        if manager_intro != '暂无':
            doc.add_heading('## 管理者介绍', level=2)
            p = doc.add_paragraph(manager_intro)
            set_chinese_font(p.runs[0])


def add_stock_analysis(doc, company_data):
    """添加股票表现分析"""
    doc.add_heading('十、股票表现分析', level=1)

    stock_items = [
        ('股票大涨时间段', company_data['股票大涨时间段']),
        ('大涨期间业绩表现', company_data['股票大涨时间段内业绩指标详细描述']),
        ('股价变动特点', company_data['此公司的股价变动特点']),
        ('是否为周期股', company_data['是否为周期股'])
    ]

    has_stock_data = False
    for label, value in stock_items:
        formatted_value = format_value(value)
        if formatted_value != '暂无':
            has_stock_data = True
            p = doc.add_paragraph()
            p.add_run(f"{label}：").bold = True
            p.add_run(formatted_value)
            set_chinese_font(p.runs[0])
            set_chinese_font(p.runs[1])

    if not has_stock_data:
        p = doc.add_paragraph('暂无股票表现分析数据')
        set_chinese_font(p.runs[0])


def add_chart_images(doc, engine, company_data):
    """添加技术分析图表 - 优化图片插入，确保4张图片最多占用2页"""
    doc.add_heading('十一、技术分析图表', level=1)

    stock_code = company_data['股票代码']

    # 获取股票图片数据
    images_df = get_stock_images_from_db(engine, stock_code)

    if images_df.empty:
        p = doc.add_paragraph('暂无技术分析图表')
        set_chinese_font(p.runs[0])
        doc.add_paragraph()  # 空行
        return

    # 按照业务关键词顺序插入图片
    images_added = 0
    for chart_name, chart_conditions in business_keywords.items():
        # 查找匹配的图片
        matching_images = []

        for chart_type, period in chart_conditions:
            # 查找匹配的记录
            for _, row in images_df.iterrows():
                if row['chart_type'] == chart_type and row['period'] == period:
                    matching_images.append(row)

        if matching_images:
            # 添加图表标题
            doc.add_heading(f'## {chart_name}', level=2)

            # 插入图片 - 使用更小的尺寸确保4张图片最多占用2页
            for image_data in matching_images:
                image_path = image_data['image_path']

                # 检查图片文件是否存在
                if os.path.exists(image_path):
                    try:
                        # 添加图片标题
                        period_name = {
                            'monthly': '月度',
                            'quarterly': '季度',
                            'yearly': '年度',
                            'all': '全周期'
                        }.get(image_data['period'], image_data['period'])

                        p = doc.add_paragraph()
                        p.alignment = WD_ALIGN_PARAGRAPH.CENTER
                        run = p.add_run(f"{period_name}{chart_name}")
                        run.italic = True
                        run.font.size = Pt(10)
                        set_chinese_font(run)

                        # 插入图片 - 使用更小的尺寸
                        # 每张图片宽度设为5英寸，高度按比例缩放
                        doc.add_picture(image_path, width=Inches(5.0))

                        # 图片后空行
                        doc.add_paragraph()
                        images_added += 1

                        # 每插入2张图片后检查是否需要分页控制
                        if images_added % 2 == 0 and images_added < len(images_df):
                            # 添加分页符，确保每页最多2张图片
                            doc.add_page_break()

                    except Exception as e:
                        print(f"插入图片失败 {image_path}: {str(e)}")
                        p = doc.add_paragraph(f"图片加载失败: {os.path.basename(image_path)}")
                        set_chinese_font(p.runs[0])
                else:
                    print(f"图片文件不存在: {image_path}")
                    p = doc.add_paragraph(f"图片文件不存在: {os.path.basename(image_path)}")
                    set_chinese_font(p.runs[0])


def add_report_footer(doc):
    """添加报告页脚"""
    p = doc.add_paragraph()
    p.alignment = WD_ALIGN_PARAGRAPH.RIGHT
    run = p.add_run(f"报告生成时间：{datetime.now().strftime('%Y年%m月%d日 %H:%M:%S')}")
    run.italic = True
    run.font.color.rgb = RGBColor(128, 128, 128)  # 灰色
    set_chinese_font(run)


def generate_company_report(company_data, engine, output_dir='company_reports'):
    """为单个公司生成Word报告"""
    try:
        # 创建输出目录
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)

        # 创建文档
        doc = Document()

        # 设置文档基本样式
        style = doc.styles['Normal']
        style.font.name = '微软雅黑'
        style._element.rPr.rFonts.set(qn('w:eastAsia'), '微软雅黑')
        style.font.size = Pt(10.5)

        # 按照要求的顺序添加各个部分
        add_company_header(doc, company_data)
        add_company_overview(doc, company_data)
        add_business_analysis(doc, company_data)
        add_market_analysis(doc, company_data)
        add_competitive_analysis(doc, company_data)
        add_financial_analysis(doc, company_data)
        add_development_analysis(doc, company_data)
        add_risk_analysis(doc, company_data)
        add_corporate_structure(doc, company_data)
        add_management_info(doc, company_data)
        add_stock_analysis(doc, company_data)

        # 添加技术分析图表
        add_chart_images(doc, engine, company_data)

        add_report_footer(doc)

        # 生成文件名并保存
        filename = f"{company_data['股票代码']}_{company_data['上市公司名称']}_详细分析报告.docx"
        filepath = os.path.join(output_dir, filename)
        doc.save(filepath)

        print(f"已生成报告：{filename}")
        return filepath

    except Exception as e:
        print(f"生成报告失败（{company_data['股票代码']}）: {str(e)}")
        return None


def batch_generate_reports(engine, limit=None):
    """批量生成所有公司报告"""
    # 读取数据
    df = read_company_detail_data(engine, limit)

    if df.empty:
        print("没有找到公司数据")
        return

    success_count = 0
    failed_companies = []

    for index, row in df.iterrows():
        print(f"正在生成 {row['股票代码']} {row['上市公司名称']} 的报告...")

        result = generate_company_report(row, engine)
        if result:
            success_count += 1
        else:
            failed_companies.append(f"{row['股票代码']} {row['上市公司名称']}")

    # 输出统计信息
    print(f"\n报告生成完成！")
    print(f"成功生成：{success_count} 份")
    print(f"失败：{len(failed_companies)} 份")

    if failed_companies:
        print("失败的公司：")
        for company in failed_companies:
            print(f"  - {company}")


if __name__ == '__main__':
    try:
        # 创建数据库连接
        engine = create_database_connection()

        # 批量生成所有公司报告
        batch_generate_reports(engine)  # 可以添加limit参数测试，如: limit=1

    except Exception as e:
        print(f"程序执行出错: {str(e)}")